import signal
import sys
import pathlib
import cv2
import rclpy
from rclpy.node import Node
from rclpy.executors import ExternalShutdownException
from cv_bridge import CvBridge
import message_filters
from sensor_msgs.msg import Image
from semantic_nav.log_utils import get_logger
logger = get_logger()
package_name = 'semantic_nav'
IMGAE_TOPIC = "/camera/color/image_raw"
DEPTH_TOPIC = "/camera/depth/image_raw"
COLOR_FRAME_HZ = 20.0
DEPTH_FRAME_HZ = 20.0


class RGBRecorder(Node):
    def __init__(self, video_path, video_max_secs=10):
        super().__init__(node_name="rgb_recorder")
        self.cv_bridge = CvBridge()
        self.video_path = video_path
        self.video_max_secs = video_max_secs
        self.file_index = 0
        self.rgb_height = None
        self.rgb_width = None
        self.declare_parameter("image_topic", IMGAE_TOPIC)
        image_topic = self.get_parameter("image_topic").value
        self.declare_parameter("color_frame_hz", COLOR_FRAME_HZ)
        self.color_frame_hz = float(self.get_parameter("color_frame_hz").value)
        self.image_subscriber = self.create_subscription(msg_type=Image, topic=image_topic, callback=self.rgb_callback, qos_profile=10)
        self.start_time = None

    def destroy_node(self):
        # Make sure to release the last set of video writers
        if self.out_rgb is not None:
            self.out_rgb.release()
        super().destroy_node()

    def initialize_video_writers(self):
        self.file_index += 1
        logger.info(f'save file index: {self.file_index}')
        rgb_filename = str((self.video_path / f'./output_rgb_{self.file_index}.mp4'))
        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.out_rgb = cv2.VideoWriter(rgb_filename, self.fourcc, self.color_frame_hz, (self.rgb_width, self.rgb_height))

    def rgb_callback(self, rgb_data):
        msg_time = rclpy.time.Time().from_msg(rgb_data.header.stamp)
        if self.rgb_height is None:
            self.start_time = msg_time
            self.rgb_height = rgb_data.height
            self.rgb_width = rgb_data.width
            self.initialize_video_writers()
        try:
            cv_image = self.cv_bridge.imgmsg_to_cv2(rgb_data, desired_encoding="bgr8")
            self.out_rgb.write(cv_image)
        except Exception as e:
            self.get_logger().error("RGB callback error: %s", e)

        # Check if time exceeds video_max_secs to switch to a new video file
        if (float((msg_time - self.start_time).nanoseconds) * 1e-9) > self.video_max_secs:
            self.start_time = msg_time  # Reset start time
            self.out_rgb.release()  # Release current VideoWriter
            self.initialize_video_writers()  # Initialize new VideoWriters for the next segment


class RGBDRecorder(Node):
    def __init__(self, video_path, video_max_secs=10):
        super().__init__(node_name="rgbd_recorder")
        self.cv_bridge = CvBridge()
        self.rgb_height = None
        self.rgb_width = None
        self.depth_width = None
        self.depth_height = None
        self.file_index = 0
        self.video_path = video_path
        self.video_max_secs = video_max_secs
        self.declare_parameter("image_topic", IMGAE_TOPIC)
        image_topic = self.get_parameter("image_topic").value
        self.declare_parameter("color_frame_hz", COLOR_FRAME_HZ)
        self.color_frame_hz = float(self.get_parameter("color_frame_hz").value)
        self.declare_parameter("depth_topic", DEPTH_TOPIC)
        depth_topic = self.get_parameter("depth_topic").value
        self.declare_parameter("depth_frame_hz", DEPTH_FRAME_HZ)
        self.depth_frame_hz = float(self.get_parameter("depth_frame_hz").value)
        image_sub = message_filters.Subscriber(self, Image, image_topic)
        depth_sub = message_filters.Subscriber(self, Image, depth_topic)
        time_syncronizer = message_filters.ApproximateTimeSynchronizer([image_sub, depth_sub], queue_size=10, slop=float(2.0 / float(max(self.color_frame_hz, self.depth_frame_hz))))
        time_syncronizer.registerCallback(self.sync_callback)
        self.start_time = None

    def destroy_node(self):
        # Make sure to release the last set of video writers
        if self.out_rgb is not None:
            self.out_rgb.release()
        if self.out_depth is not None:
            self.out_rgb.release()
        super().destroy_node()

    def initialize_video_writers(self):
        self.file_index += 1
        logger.info(f'save file index: {self.file_index}')
        rgb_filename = str((self.video_path / f'./output_rgb_{self.file_index}.mp4'))
        depth_filename = str((self.video_path / f'./output_depth_{self.file_index}.mp4'))

        self.fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.out_rgb = cv2.VideoWriter(rgb_filename, self.fourcc, self.color_frame_hz, (self.rgb_width, self.rgb_height))
        self.out_depth = cv2.VideoWriter(depth_filename, self.fourcc, self.depth_frame_hz, (self.depth_width, self.depth_height))

    def sync_callback(self, rgb_data, depth_data):
        msg_time = rclpy.time.Time().from_msg(rgb_data.header.stamp)
        if self.rgb_height is None:
            self.start_time = msg_time
            self.rgb_height = rgb_data.height
            self.rgb_width = rgb_data.width
            self.depth_width = depth_data.width
            self.depth_height = depth_data.height
            self.initialize_video_writers()
        try:
            cv_rgb_image = self.cv_bridge.imgmsg_to_cv2(rgb_data, desired_encoding="bgr8")
            self.out_rgb.write(cv_rgb_image)
            cv_depth_image = self.cv_bridge.imgmsg_to_cv2(depth_data, desired_encoding="32FC1")
            depth_colored = cv2.applyColorMap(cv2.convertScaleAbs(cv_depth_image, alpha=0.03), cv2.COLORMAP_JET)
            self.out_depth.write(depth_colored)
        except Exception as e:
            self.get_logger().error("Sync callback error: %s", e)

        # Check if time exceeds video_max_secs to switch to a new video file
        if (float((msg_time - self.start_time).nanoseconds) * 1e-9) > self.video_max_secs:
            self.start_time = msg_time  # Reset start time
            self.out_rgb.release()  # Release current VideoWriter
            self.out_depth.release()  # Release current VideoWriter
            self.initialize_video_writers()  # Initialize new VideoWriters for the next segment


def main(args=None):
    rclpy.init(args=args)
    try:
        from ament_index_python.packages import get_package_prefix
        package_path = get_package_prefix(package_name)
        video_path = pathlib.Path(package_path).joinpath('video')
        video_path.mkdir(parents=True, exist_ok=True)
        # Only read image data
        recorder = RGBRecorder(video_path)

        # Uncomment the following code for synchronized reading of image and depth data
        # recorder = RGBDRecorder(video_path)
        rclpy.spin(recorder)
        recorder.destroy_node()
    except KeyboardInterrupt:
        pass
    except ExternalShutdownException:
        pass
    finally:
        recorder.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
